"""
器官分割与重建流程第三步，从Nifti掩码到VTP三维模型，并且只关注指定器官
"""


import os
import vtk
from typing import Dict, Optional, List,Any

class ReconstructionService:
    def __init__(self, input_dir: str, output_dir: str):
        """
        初始化重建服务。

        :param input_dir: 存放100多个.nii.gz器官掩码文件的目录。
        :param output_dir: 用于保存生成的.vtp模型文件的目录。
        """
        if not os.path.isdir(input_dir):
            raise FileNotFoundError(f"输入目录不存在或不是一个有效目录: {input_dir}")
            
        os.makedirs(output_dir, exist_ok=True)  
        self.input_dir = input_dir
        self.output_dir = output_dir
        print(f"ReconstructionService 已初始化。")
        print(f"  - Nifti 输入源: {self.input_dir}")   
        print(f"  - VTP 输出目标: {self.output_dir}")

    def _get_reconstruction_params(self,organ_name:str) -> Dict[str, Any]:
        """
        获取特定器官的重建参数。

        :param organ_name: 器官名称
        :return: 包含重建参数的字典,针对器官返回最优的重建参数
        """

        params={'smoothing':100,'decimation':0.9,'pass_band':0.005}
        #针对不同器官设置不同重建参数
        # if any(kw in organ_name for kw in ['aorta','artery','vein','vena_cava','portal','portal_vein','duedenum','esophagus']):
        #     params['smoothing'] = 20
        #     params['decimation'] = 0.3
        #     params['pass_band'] = 0.05

        # elif any(kw in organ_name for kw in ['liver','spleen','gallbladder','kidney','lung']):
        #     params['smoothing'] = 100
        #     params['decimation'] = 0.9
        #     params['pass_band'] = 0.005
        # elif any(kw in organ_name for kw in ['vertebrae','rib','hib','sacrum','sternum']):
        #     params['smoothing'] = 75
        #     params['decimation'] = 0.5
        #     params['pass_band'] = 0.01
        return params

    def reconstruct(self,organs_of_interest:Optional[List[str]]=None,mask_value:int=1) -> List[str]:
        """
        执行批量重建的核心公共方法。

        :param organs_of_interest: 可选参数，指定需要重建的器官列表。如果为 None，则处理所有器官。
        :return: 一个包含了所有成功生成的VTP文件名的列表。
        """
        print("\n--- 开始执行批量重建任务 ---")
        
        all_files = [f for f in os.listdir(self.input_dir) if f.endswith(('.nii', '.nii.gz'))]
        
        if organs_of_interest:
            interest_set = set(organs_of_interest)
            files_to_process = [f for f in all_files if f.split('.')[0] in interest_set]
            print(f"将处理指定的 {len(files_to_process)} 个器官。")
        else:
            files_to_process = all_files
            print(f"将处理所有 {len(files_to_process)} 个器官。")

        successfully_generated = []
        for i, filename in enumerate(files_to_process):
            input_path = os.path.join(self.input_dir, filename)
            organ_name = filename.split('.')[0]
            output_path = os.path.join(self.output_dir, f"{organ_name}.vtp")
            
            params = self._get_reconstruction_params(organ_name)
            print(f"\n[{i+1}/{len(files_to_process)}] 处理: {organ_name} (平滑:{params['smoothing']}, 简化:{params['decimation']})")

            try:
                reader = vtk.vtkNIFTIImageReader()
                reader.SetFileName(input_path)          # 设置NIFTI文件路径

                # 使用 vtkDiscreteMarchingCubes
                mc = vtk.vtkDiscreteMarchingCubes()
                mc.SetInputConnection(reader.GetOutputPort())
                mc.GenerateValues(mask_value, mask_value, mask_value)
                mc.Update()

                if mc.GetOutput().GetNumberOfPoints() == 0:
                    print(f"  - 警告: '{organ_name}' 模型为空，已跳过。")
                    continue
                
                # 连接到平滑器
                smoother = vtk.vtkWindowedSincPolyDataFilter()
                smoother.SetInputConnection(mc.GetOutputPort())
                smoother.SetNumberOfIterations(params['smoothing'])
                smoother.SetPassBand(params['pass_band'])
                smoother.Update()

                # 连接到简化器
                decimator = vtk.vtkDecimatePro()
                decimator.SetInputConnection(smoother.GetOutputPort())
                decimator.SetTargetReduction(params['decimation'])
                decimator.PreserveTopologyOn()
                decimator.Update()

                # 写入VTP文件
                writer = vtk.vtkXMLPolyDataWriter()
                writer.SetFileName(output_path)
                writer.SetInputConnection(decimator.GetOutputPort())
                writer.Write()
                
                print(f"  - 成功生成: {os.path.basename(output_path)}")
                successfully_generated.append(os.path.basename(output_path))
            except Exception as e:
                print(f"  - 错误: 处理 '{filename}' 时发生错误: {e}")
        
        print(f"\n🎉 --- 批量重建任务完成！共成功生成 {len(successfully_generated)} 个模型。--- 🎉")
        return successfully_generated

if __name__ == '__main__':
    BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    NIFTI_MASKS_DIR = os.path.join(BASE_DIR, r"data\segmentation_results\\patient001\\Abdomen-V  2.0  B30f")
    VTP_MODELS_OUTPUT_DIR = os.path.join(BASE_DIR, r"data\\vtp_models\\patient001\\Abdomen-V  2.0  B30f\\all_organs")
    # eus_organs_of_interest = [
    #     # --- 消化道 (EUS探头路径和观察窗口) ---
    #     'esophagus', 
    #     'stomach', 
    #     'duodenum',

    #     # --- 核心目标器官 ---
    #     'pancreas',
    #     'liver',
    #     'gallbladder',
    #     'spleen',
    #     'kidney_left',
    #     'kidney_right',
    #     'adrenal_gland_left',   # 左肾上腺
    #     'adrenal_gland_right',  # 右肾上腺

    #     # 大血管主干
    #     'aorta',                      # 腹主动脉
    #     'inferior_vena_cava',       # 下腔静脉
    
    #     # 腹腔干及其分支 (Celiac Trunk)
    #     'celiac_trunk',               # 腹腔干
    #     'splenic_artery',             # 脾动脉
    #     # 'hepatic_artery_proper',    # 肝固有动脉 (TotalSegmentator v2 新增)

    #     # 肠系膜血管
    #     'superior_mesenteric_artery', # 肠系膜上动脉
    #     'superior_mesenteric_vein',   # 肠系膜上静脉

    #     # 门静脉系统
    #     'portal_vein_and_splenic_vein', # 门静脉和脾静脉 (TotalSegmentator将它们合并)
    #     'hepatic_vein',               # 肝静脉

    #     # 肾血管
    #     'renal_artery_left',          # 左肾动脉
    #     'renal_artery_right',         # 右肾动脉
    #     'renal_vein_left',            # 左肾静脉
    #     'renal_vein_right',           # 右肾静脉
    
    #     # --- 骨骼参照 (用于空间定位) ---
    #     'vertebrae_L1', 
    #     'vertebrae_L2'
    # ]

    try:
        service = ReconstructionService(NIFTI_MASKS_DIR, VTP_MODELS_OUTPUT_DIR)
        # 只重建我们感兴趣的器官
        service.reconstruct()
    except FileNotFoundError as e:
        print(f"\n错误: {e}")